import type { Ref } from '../base'

export class RefIndex<E, T> {
  private data: Record<Ref<E>, Set<T>> = {}

  constructor(lst: T[], fn: (t: T) => Ref<E>[] | Ref<E> | undefined | null) {
    for (const t of lst) {
      const r = fn(t)
      if (!r) {
        continue
      }
      if (r instanceof Array) {
        for (const rr of r) {
          let s = this.data[rr]
          if (!s) {
            s = new Set<T>()
            this.data[rr] = s
          }
          s.add(t)
        }
        continue
      }
      let s = this.data[r]
      if (!s) {
        s = new Set<T>()
        this.data[r] = s
      }
      s.add(t)
    }
  }

  public get(e: Ref<E>): Set<T> | undefined {
    return this.data[e]
  }
}
